﻿Option Strict On

Imports System.Management


''' <summary>
''' A disk description + a collection of functions to work with disks via WMI
''' </summary>
''' <remarks></remarks>
Public Class WmiDisk

    ''' <summary>
    ''' Drive index, 0-based
    ''' </summary>
    ''' <remarks></remarks>
    Public Index As Integer = 0

    ''' <summary>
    ''' True if the disk a fixed hard drive
    ''' </summary>
    ''' <remarks></remarks>
    Public IsHDD As Boolean = False

    ''' <summary>
    ''' Number or partitions
    ''' </summary>
    ''' <remarks></remarks>
    Public Partitions As Integer = 0

    ''' <summary>
    ''' Set to true when the disk was accounted for
    ''' </summary>
    ''' <remarks></remarks>
    Public Processed As Boolean = False

    ''' <summary>
    ''' Letter of the disk (one per disk only)
    ''' </summary>
    ''' <remarks></remarks>
    Public Letter As String = Nothing


    Public Sub New(ByRef oProps As Management.PropertyDataCollection)

        For Each oProp As Management.PropertyData In oProps


            Dim sName As String = oProp.Name
            Select Case sName

                Case "Index"
                    If Not oProp.Value Is Nothing Then Index = CInt(oProp.Value)
                Case "MediaType"
                    If Not oProp.Value Is Nothing AndAlso oProp.Value.ToString.ToLower = "fixed hard disk media" Then IsHDD = True
                Case "Partitions"
                    If Not oProp.Value Is Nothing Then Partitions = CInt(oProp.Value)
            End Select
        Next




    End Sub



    ''' <summary>
    ''' Rescan drives on the machines and load them into arDisksPhys and returns the number of drives
    ''' </summary>
    ''' <remarks></remarks>
    Public Shared Function WmiLoadDrives() As Integer

        'reset list of disks
        EnviroBuilder.arDisksPhys = New List(Of WmiDisk)

        'prepare scope
        Dim scopeStr As String = "\\.\root\cimv2"
        Dim scope As New ManagementScope(scopeStr)
        scope.Connect()

        'prepare WMQ
        Dim queryString As String = "SELECT * FROM Win32_DiskDrive"  'Win32_Volume WHERE DriveLetter IS NOT NULL" Win32_DiskDrive
        Dim query As New SelectQuery(queryString)

        'query and iterate
        Using searcher As New ManagementObjectSearcher(scope, query)
            For Each disk As ManagementObject In searcher.Get()
                'create a WmiDisk class and add it to the list
                Dim oProps As Management.PropertyDataCollection = disk.Properties
                EnviroBuilder.arDisksPhys.Add(New WmiDisk(oProps))
            Next
        End Using

        'return number of disks
        Return EnviroBuilder.arDisksPhys.Count

    End Function


    ''' <summary>
    ''' Formats the first unformatted volume with default settings.
    ''' </summary>
    ''' <remarks></remarks>
    Public Shared Function WmiFormatVolumes() As Boolean

        'volume name cannot be nothing
        'If VolumeName Is Nothing Then VolumeName = ""

        'connect to WMI
        Dim scopeStr As String = "\\.\root\cimv2"
        Dim scope As New ManagementScope(scopeStr)
        scope.Connect()

        'prepare WMQ
        Dim queryString As String = "SELECT * FROM Win32_Volume"  'Win32_Volume WHERE DriveLetter IS NOT NULL" Win32_DiskDrive
        Dim query As New SelectQuery(queryString)

        'search and  iterate through results
        Using searcher As New ManagementObjectSearcher(scope, query)

            For Each disk As ManagementObject In searcher.Get()

                'get disk properties
                Dim oProps As Management.PropertyDataCollection = disk.Properties

                'check the file system
                Dim sFS As String = Nothing
                Dim oFS As Management.PropertyData = oProps("FileSystem")
                If Not oFS.Value Is Nothing Then sFS = oProps("FileSystem").Value.ToString()
                'format the drive if needed
                If sFS = "" Then
                    Dim args(4) As Object
                    args(0) = CStr("NTFS")
                    args(1) = True
                    args(2) = CType(4096, UInteger)
                    args(3) = CStr("")
                    args(4) = False
                    Dim iRes As Integer = -1
                    Try
                        iRes = CInt(disk.InvokeMethod("Format", args))
                    Catch ex As Exception
                        EnviroBuilder.WriteToLog("Cannot format volume. " + vbCrLf + args.ToString + vbCrLf + ex.Message, EventLogEntryType.Error)
                        Return False
                    End Try

                    'formatted OK?
                    If iRes <> 0 Then
                        EnviroBuilder.WriteToLog("Format volume failed with error code " + iRes.ToString, EventLogEntryType.Error)
                        Return False

                    End If

                End If
            Next
        End Using

        'nothing to format if reached this point
        Return True

    End Function

    ''' <summary>
    ''' Retuns first 0-based disk number where no drive letter assigned or -1 for none.
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Shared Function WmiUnnamedVolumeDiskIndex() As Integer

        'volume name cannot be nothing
        'If VolumeName Is Nothing Then VolumeName = ""

        Dim sLog As String = "Looking for an unnamed volume." + vbCrLf

        'connect to WMI
        Dim scopeStr As String = "\\.\root\cimv2"
        Dim scope As New ManagementScope(scopeStr)
        scope.Connect()

        'prepare WMQ
        Dim queryString As String = "SELECT * FROM Win32_LogicalDiskToPartition"
        Dim query As New SelectQuery(queryString)

        'search and  iterate through results
        Using searcher As New ManagementObjectSearcher(scope, query)
            For Each disk As ManagementObject In searcher.Get()
                'get disk number from a string like this \\IP-3ABBE52A\root\cimv2:Win32_DiskPartition.DeviceID="Disk #4, Partition #0"
                Dim oProps As Management.PropertyDataCollection = disk.Properties
                'get disk ID
                Dim sFS As String = oProps("Antecedent").Value.ToString
                Dim iDiskIdx As Integer = CInt(sFS.Substring(sFS.IndexOf("#"c) + 1, 1))
                sLog += vbCrLf + sFS + vbCrLf + iDiskIdx.ToString + vbCrLf
                'get disk letter
                sFS = oProps("Dependent").Value.ToString
                Dim sLetter As String = sFS.Substring(sFS.Length - 3, 1)
                sLog += sFS + vbCrLf + sLetter + vbCrLf
                'find a matching disk and set the letter
                For Each oD As WmiDisk In EnviroBuilder.arDisksPhys
                    If oD.Index = iDiskIdx Then
                        oD.Letter = sLetter
                        sLog += "Disk mapped" + vbCrLf
                        Exit For
                    End If
                Next
            Next
        End Using

        'select disk indexes with no letters
        Dim arDisks As New List(Of Integer)
        For Each oD As WmiDisk In EnviroBuilder.arDisksPhys
            If oD.IsHDD AndAlso oD.Partitions = 1 AndAlso oD.Letter = "" Then arDisks.Add(oD.Index)
        Next

        'return nothing if no disks need letter assignment
        If arDisks.Count = 0 Then
            sLog += "No disks without letters found"
            EnviroBuilder.WriteToLog(sLog, EventLogEntryType.Information)
            Return -1
        End If

        'return the first one without a letter
        arDisks.Sort()
        Dim iOut As Integer = arDisks(0)
        sLog += "Disk without a letter: " + iOut.ToString
        EnviroBuilder.WriteToLog(sLog, EventLogEntryType.Information)

        Return iOut

    End Function


End Class
